#!/usr/bin/env python
#coding:utf-8

# 手势追踪
import rospy
import cv2
from geometry_msgs.msg import Twist
from sensor_msgs.msg import CompressedImage
from cv_bridge import CvBridge, CvBridgeError
import numpy as np
import mediapipe as mp
import math



# ------------------------------计算角度函数------------------------------
def vector_2d_angle(v1, v2):
    v1_x = v1[0]
    v1_y = v1[1]
    v2_x = v2[0]
    v2_y = v2[1]
    try:
        angle_ = math.degrees(math.acos(
        (v1_x * v2_x + v1_y * v2_y) / (((v1_x ** 2 + v1_y ** 2) ** 0.5) * ((v2_x ** 2 + v2_y ** 2) ** 0.5))))
    except:
        angle_ = 65535.
        if angle_ > 180.:
            angle_ = 65535.
    return angle_


# ------------------------------计算每个手指弯曲角度函数------------------------------
def hand_angle(hand_):
    angle_list = []
    # ---------------------------- thumb 大拇指角度
    angle_ = vector_2d_angle(
    ((int(hand_[0][0]) - int(hand_[2][0])), (int(hand_[0][1]) - int(hand_[2][1]))),
    ((int(hand_[3][0]) - int(hand_[4][0])), (int(hand_[3][1]) - int(hand_[4][1])))
    )
    angle_list.append(angle_)
    # ---------------------------- index 食指角度
    angle_ = vector_2d_angle(
    ((int(hand_[0][0]) - int(hand_[6][0])), (int(hand_[0][1]) - int(hand_[6][1]))),
    ((int(hand_[7][0]) - int(hand_[8][0])), (int(hand_[7][1]) - int(hand_[8][1])))
    )
    angle_list.append(angle_)
    # ---------------------------- middle 中指角度
    angle_ = vector_2d_angle(
    ((int(hand_[0][0]) - int(hand_[10][0])), (int(hand_[0][1]) - int(hand_[10][1]))),
    ((int(hand_[11][0]) - int(hand_[12][0])), (int(hand_[11][1]) - int(hand_[12][1])))
    )
    angle_list.append(angle_)
    # ---------------------------- ring 无名指角度
    angle_ = vector_2d_angle(
    ((int(hand_[0][0]) - int(hand_[14][0])), (int(hand_[0][1]) - int(hand_[14][1]))),
    ((int(hand_[15][0]) - int(hand_[16][0])), (int(hand_[15][1]) - int(hand_[16][1])))
    )
    angle_list.append(angle_)
    # ---------------------------- pink 小拇指角度
    angle_ = vector_2d_angle(
    ((int(hand_[0][0]) - int(hand_[18][0])), (int(hand_[0][1]) - int(hand_[18][1]))),
    ((int(hand_[19][0]) - int(hand_[20][0])), (int(hand_[19][1]) - int(hand_[20][1])))
    )
    angle_list.append(angle_)
    return angle_list


# ------------------------------判断手指指令函数------------------------------
def h_gesture(angle_list):
    thr_angle = 65.
    thr_angle_thumb = 53.
    thr_angle_s = 49.
    gesture_str = None
    index = 6
    if 65535. not in angle_list:
        if (angle_list[0] > thr_angle_thumb) and (angle_list[1] > thr_angle) and (angle_list[2] > thr_angle) and (
        angle_list[3] > thr_angle) and (angle_list[4] > thr_angle):
            gesture_str = "off"
            index = 5
        elif (angle_list[0] < thr_angle_s) and (angle_list[1] < thr_angle_s) and (angle_list[2] < thr_angle_s) and (
        angle_list[3] < thr_angle_s) and (angle_list[4] < thr_angle_s):
            gesture_str = "on"
            index = 6
        elif (angle_list[0] < thr_angle_s) and (angle_list[1] > thr_angle) and (angle_list[2] > thr_angle) and (         
                    angle_list[3] > thr_angle) and (angle_list[4] > thr_angle):         
            gesture_str = "thumbUp"
            index = 7
    return gesture_str, index


# ------------------------------判断手在图像中的偏移位置函数------------------------------
def location(hand_):
    x = hand_[9][0] - 320
    y = hand_[9][1] -240
    return x, y


def main():
    # ROS节点初始化
    rospy.init_node('gesture_publisher', anonymous=True)

    # 创建两个Publisher，发布名为/vision/camera/image_raw的topic，消息类型为sensor_msgs::Image，队列长度10
    img_raw_pub = rospy.Publisher('/vision/camera/image_raw',CompressedImage, queue_size=10)
    img_hd_pub = rospy.Publisher('/vision/camera/image_handle',CompressedImage, queue_size=10)
    # 创建一个Publisher，发布名为/gesture/cmd_vel的topic，消息类型为geometry_msgs::Twist，队列长度10
    gesture_vel_pub = rospy.Publisher('/gesture/cmd_vel', Twist, queue_size=10)

    bridge = CvBridge()

    # 创建压缩图像
    msg = CompressedImage()

    mp_drawing = mp.solutions.drawing_utils
    mp_hands = mp.solutions.hands
    hands = mp_hands.Hands(
    static_image_mode=False,
    max_num_hands=1,
    min_detection_confidence=0.75,
    min_tracking_confidence=0.75)
    cap = cv2.VideoCapture(0)
    ges = ''
    while cap.isOpened():

        ret, frame = cap.read()
        frame = cv2.flip(frame, 1)
        # 发布原始图像
        msg.header.stamp = rospy.Time.now()
        msg.format = "jpeg"
        msg.data = np.array(cv2.imencode('.jpg', frame)[1]).tostring()
        img_raw_pub.publish(msg)

        # print(np.shape(frame)) # 图像大小
        frame2 = frame
        frame2 = cv2.cvtColor(frame2, cv2.COLOR_BGR2RGB)
        results = hands.process(frame2)
        frame2 = cv2.cvtColor(frame2, cv2.COLOR_RGB2BGR)

        if results.multi_hand_landmarks:
            for hand_landmarks in results.multi_hand_landmarks:
                # print('hand_landmarks:', hand_landmarks)
                # 21个关键点可视化
                mp_drawing.draw_landmarks(frame2, hand_landmarks, mp_hands.HAND_CONNECTIONS)
                hand_local = []
                for i in range(21):
                    x = hand_landmarks.landmark[i].x * frame.shape[1]
                    y = hand_landmarks.landmark[i].y * frame.shape[0]
                    hand_local.append((x, y))
                if hand_local:
                    angle_list = hand_angle(hand_local)  # 计算每个手指的弯曲角度
                    gesture_str, ind = h_gesture(angle_list)  # 判断指令
                    if gesture_str != ges:
                        # print(gesture_str)
                        ges = gesture_str
                    cv2.putText(frame2, gesture_str, (0, 100), 0, 1.3, (0, 0, 255), 3)
                    x, y = location(hand_local)    # 判断手在图像中的偏移位置

                    # 初始化geometry_msgs::Twist类型的消息
                    vel_msg = Twist()
                    vel_msg.linear.x = float(x)
                    vel_msg.linear.y = float(y)
                    vel_msg.linear.z = float(ind)
                    # 发布消息
                    gesture_vel_pub.publish(vel_msg)
                    rospy.loginfo("%0.2f, %0.2f", vel_msg.linear.x, vel_msg.linear.y)

        # 发布处理后的图像
        msg.header.stamp = rospy.Time.now()
        msg.format = "jpeg"
        msg.data = np.array(cv2.imencode('.jpg', frame2)[1]).tostring()
        img_hd_pub.publish(msg)



if __name__ == "__main__":
    try:
        main()
    except rospy.ROSInterruptException:
        pass


